#include "Thread.h"

NTSTATUS Thread::ApcpQuerySystemProcessInformation(PSYSTEM_PROCESS_INFO * SystemInfo)
{
	PSYSTEM_PROCESS_INFO pBuffer = NULL;
	ULONG BufferSize = 0;
	ULONG RequiredSize = 0;

	NTSTATUS status = STATUS_SUCCESS;
	while ((status = ZwQuerySystemInformation(
		SystemProcessInformation,
		pBuffer,
		BufferSize,
		&RequiredSize//retn Length
	)) == STATUS_INFO_LENGTH_MISMATCH)
	{
		BufferSize = RequiredSize;
		pBuffer = (PSYSTEM_PROCESS_INFO)ExAllocatePool(PagedPool, BufferSize);
	}

	if (!NT_SUCCESS(status))
	{
		if (pBuffer != NULL)
		{
			ExFreePool(pBuffer);
		}

		return status;
	}
	//retn pSystemProcessInfo
	*SystemInfo = pBuffer;
	return status;
}

PVOID Thread::GetWin32ThreadAddress(PEPROCESS pEProcess)
{
	PSYSTEM_PROCESS_INFO OriginalSystemProcessInfo = NULL;

	NTSTATUS status = ApcpQuerySystemProcessInformation(&OriginalSystemProcessInfo);

	PVOID Win32Thread = NULL;

	PETHREAD Pethread = NULL;

	if (!NT_SUCCESS(status))
	{
		ObDereferenceObject(pEProcess);
		return NULL;
	}


	PSYSTEM_PROCESS_INFO SystemProcessInfo = OriginalSystemProcessInfo;
	status = STATUS_NOT_FOUND;
	do
	{
		if (SystemProcessInfo->UniqueProcessId == PsGetProcessId(pEProcess))
		{
			status = STATUS_SUCCESS;
			break;
		}

		SystemProcessInfo = (PSYSTEM_PROCESS_INFO)((PUCHAR)SystemProcessInfo + SystemProcessInfo->NextEntryOffset);
	} while (SystemProcessInfo->NextEntryOffset != 0);


	if (!NT_SUCCESS(status))
	{
		ExFreePool(OriginalSystemProcessInfo);
		return NULL;
	}
	
	for (ULONG Index = 0; Index < SystemProcessInfo->NumberOfThreads; ++Index){

		
		status = PsLookupThreadByThreadId(SystemProcessInfo->Threads[Index].ClientId.UniqueThread, &Pethread);
		if (!NT_SUCCESS(status))
		{
			continue;
		}
		Win32Thread = PsGetThreadWin32Thread(Pethread);
		ObDereferenceObject(Pethread);
		if (Win32Thread != NULL) {
			break;
		}
	}
	
	return Win32Thread;
}
